package com.lw.leetcode.arr.b;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 973. 最接近原点的 K 个点
 *
 * @Author liw
 * @Date 2021/8/31 22:13
 * @Version 1.0
 */
public class KClosest {

    public static void main(String[] args) {
        KClosest test = new KClosest();
        PriorityQueue<Long> queue = new PriorityQueue<>((a, b) -> Long.compare(b, a));
        queue.add(1L);
        queue.add(3L);
        queue.add(2L);
        System.out.println(queue.poll());
        System.out.println(queue.poll());
        System.out.println(queue.poll());
    }

    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> Long.compare(b[0] * b[0] + b[1] * b[1], a[0] * a[0] + a[1] * a[1]));
        for (int i = 0; i < k; i++) {
            queue.add(points[i]);
        }
        int length = points.length;
        for (int i = k; i < length; i++) {
            queue.add(points[i]);
            queue.poll();
        }
        int[][] arr = new int[k][2];
        for (int i = 0; i < k; i++) {
            arr[i] = queue.poll();
        }
        return arr;
    }

    public int[][] kClosest2(int[][] points, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(new Comparator<int[]>() {
            public int compare(int[] array1, int[] array2) {
                return array2[0] - array1[0];
            }
        });
        for (int i = 0; i < k; ++i) {
            pq.offer(new int[]{points[i][0] * points[i][0] + points[i][1] * points[i][1], i});
        }
        int n = points.length;
        for (int i = k; i < n; ++i) {
            int dist = points[i][0] * points[i][0] + points[i][1] * points[i][1];
            if (dist < pq.peek()[0]) {
                pq.poll();
                pq.offer(new int[]{dist, i});
            }
        }
        int[][] ans = new int[k][2];
        for (int i = 0; i < k; ++i) {
            ans[i] = points[pq.poll()[1]];
        }
        return ans;
    }

}
